# -*- coding: utf-8 -*-
# $ ./build/triton ./src/examples/pin/inject_model_with_snapshot.py ./src/samples/crackmes/crackme_xor a
from triton import *
from pintool import *
import sys

Triton = getTritonContext()

def csym(instruction):
    if instruction.getAddress() == 0x8048CF4:
        esp = getCurrentRegisterValue(Triton.registers.esp)
        Triton.convertMemoryToSymbolicVariable(MemoryAccess(esp, CPUSIZE.DWORD))
        return
    return


def cafter(instruction):
    # 400597 cmp     ecx, eax
    # 开始求解约束
    if instruction.getAddress() == 0X8048CFC:
        astCtxt = Triton.getAstContext()
        zf = Triton.getRegisterAst(Triton.registers.zf)
        # 如果正确的话，需要 zf == 1， 增加约束
        cstr = astCtxt.land([
            Triton.getPathConstraintsAst(),
            zf == 1
        ])
        print(cstr)
        models = Triton.getModel(cstr)
        for k, v in list(models.items()):
            # 把计算的结果保存
            print(v.getValue())
        return
    return


def fini():
    print('[+] Analysis done!')
    return




if __name__ == '__main__':
    setupImageWhitelist(['bomb'])
    # 设置从 check 函数开始分析
    # startAnalysisFromSymbol('check')
    startAnalysisFromAddress(0x8048CB1)
    
    insertCall(cafter,  INSERT_POINT.AFTER)
    insertCall(csym,    INSERT_POINT.BEFORE_SYMPROC)

    insertCall(fini,    INSERT_POINT.FINI)
    runProgram()